package fem2d;

import inf.jlinalg.SolveFailedException;
import inf.math.LinAlg;
import inf.text.ArrayFormat;

import java.io.IOException;
import java.util.ArrayList;

import junit.framework.Assert;
import math2.ConstantFunctionRnToR;
import math2.FunctionRnToR;

import org.junit.Test;

import fem2.Constraint;
import fem2.DistributedLoad;
import fem2.Edge;
import fem2.Element;
import fem2.Face;
import fem2.Load;
import fem2.MaterialModel;
import fem2.Mesh;
import fem2.MeshUtilities;
import fem2.Model;
import fem2.Node;
import fem2.analysis.Analysis;
import fem2.analysis.LinearStaticAnalysis;
import fem2.element.BBarStructuralElement;
import fem2.element.EASStructuralElement;
import fem2.element.StructuralElement;
import fem2.enu.EASModeType;
import fem2.enu.EchoLevelType;
import fem2.enu.IntegrationType;
import fem2.enu.SolverType;
import fem2.enu.State;
import fem2.enu.StressType;
import fem2.jlinalg.BaitschSolver;
import fem2.jlinalg.Solver;
import fem2.material.StVenantKirchhoffMM;
import fem2.observer.StressObserver;
import fem2.pre_and_post.GidPostStructural;

/**
 * desmonstrate volumetric locking demo with B-Bar method, selective reduced
 * integration, ...
 * 
 * @author hbui
 * 
 */
public class volumetric_locking_demo_quadratic {

	public Model createModel(int type) {
		Mesh mesh = new Mesh();

		Node n1 = mesh.addNode(new Node(0, 1));
		Node n2 = mesh.addNode(new Node(0.2, 0.98));
		Node n3 = mesh.addNode(new Node(0.38, 0.92));
		Node n4 = mesh.addNode(new Node(0.56, 0.83));
		Node n5 = mesh.addNode(new Node(0.71, 0.71));
		Node n6 = mesh.addNode(new Node(0.83, 0.56));
		Node n7 = mesh.addNode(new Node(0.92, 0.38));
		Node n8 = mesh.addNode(new Node(0.98, 0.2));
		Node n9 = mesh.addNode(new Node(1, 0));
		Node n10 = mesh.addNode(new Node(0, 1.75));
		Node n11 = mesh.addNode(new Node(0.67, 1.62));
		Node n12 = mesh.addNode(new Node(1.24, 1.24));
		Node n13 = mesh.addNode(new Node(1.62, 0.67));
		Node n14 = mesh.addNode(new Node(1.75, 0));
		Node n15 = mesh.addNode(new Node(0, 2.5));
		Node n16 = mesh.addNode(new Node(0.49, 2.45));
		Node n17 = mesh.addNode(new Node(0.96, 2.31));
		Node n18 = mesh.addNode(new Node(1.39, 2.08));
		Node n19 = mesh.addNode(new Node(1.77, 1.77));
		Node n20 = mesh.addNode(new Node(2.08, 1.39));
		Node n21 = mesh.addNode(new Node(2.31, 0.96));
		Node n22 = mesh.addNode(new Node(2.45, 0.49));
		Node n23 = mesh.addNode(new Node(2.5, 0));
		Node n24 = mesh.addNode(new Node(0, 3.25));
		Node n25 = mesh.addNode(new Node(1.24, 3));
		Node n26 = mesh.addNode(new Node(2.3, 2.3));
		Node n27 = mesh.addNode(new Node(3, 1.24));
		Node n28 = mesh.addNode(new Node(3.25, 0));
		Node n29 = mesh.addNode(new Node(0, 4));
		Node n30 = mesh.addNode(new Node(0.78, 3.92));
		Node n31 = mesh.addNode(new Node(1.53, 3.7));
		Node n32 = mesh.addNode(new Node(2.22, 3.33));
		Node n33 = mesh.addNode(new Node(2.83, 2.83));
		Node n34 = mesh.addNode(new Node(3.33, 2.22));
		Node n35 = mesh.addNode(new Node(3.7, 1.53));
		Node n36 = mesh.addNode(new Node(3.92, 0.78));
		Node n37 = mesh.addNode(new Node(4, 0));

		mesh.addFace(new Face(n1, n3, n17, n15, n2, n11, n16, n10));
		mesh.addFace(new Face(n3, n5, n19, n17, n4, n12, n18, n11));
		mesh.addFace(new Face(n5, n7, n21, n19, n6, n13, n20, n12));
		mesh.addFace(new Face(n7, n9, n23, n21, n8, n14, n22, n13));
		mesh.addFace(new Face(n15, n17, n31, n29, n16, n25, n30, n24));
		mesh.addFace(new Face(n17, n19, n33, n31, n18, n26, n32, n25));
		mesh.addFace(new Face(n19, n21, n35, n33, n20, n27, n34, n26));
		mesh.addFace(new Face(n21, n23, n37, n35, n22, n28, n36, n27));

		// String projDir = "/home/hbui/workspace2/src-matlab/locking";
		// String projName = "";
		// String fn = projDir + '/' + projName + '/' +
		// "volumetric_locking_demo_quadratic.msh";
		// MeshGenerator mg = new GidMeshGenerator(fn);
		// Mesh mesh = mg.getMesh(2);

		Model m = new Model(mesh);

		/*
		 * element creation
		 */
		double mu = 1;
		double nu = 0.499;
		double E = 2 * mu * (1 + nu);
		double t = 1.0;
		State ss = State.PLANE_STRAIN;
		MaterialModel mm = new StVenantKirchhoffMM(E, nu, t, 0, ss);
		for (int i = 0; i < mesh.countFaces(); i++) {
			StructuralElement e;
			if (type == 0) {
				e = new StructuralElement(mesh.getFace(i), mm);
			} else if (type == 1) {
				e = new BBarStructuralElement(mesh.getFace(i), mm);
			} else if (type == 2) {
				e = new EASStructuralElement(mesh.getFace(i), mm, EASModeType.Q1E4);
			} else if (type == 3) {
				e = new EASStructuralElement(mesh.getFace(i), mm, EASModeType.Q1E5);
			} else {
				throw new Error("invalid element type");
			}
			m.addElement(e);
		}

		/*
		 * conditions
		 */
		Constraint c1 = new Constraint(false, true);
		Constraint c2 = new Constraint(true, false);

		ArrayList<Node> leftedge = MeshUtilities.seekNodesOnSurface(mesh, 1.0, 0.0, 0.0);
		ArrayList<Node> botedge = MeshUtilities.seekNodesOnSurface(mesh, 0.0, 1.0, 0.0);

		mesh.setConstraint(c1, leftedge);
		mesh.setConstraint(c2, botedge);

		double[][] traction = new double[][] { { 0.195090322, 0.98078528 },
				{ 0.555570233, 0.831469612 }, { 0.831469612, 0.555570233 },
				{ 0.98078528, 0.195090322 } };
		int[][] tractionedges = new int[][] { { 1, 2, 3 }, { 3, 4, 5 }, { 5, 6, 7 }, { 7, 8, 9 } };

		for (int i = 0; i < tractionedges.length; i++) {
			int[] edge = tractionedges[i];
			FunctionRnToR ldx = new ConstantFunctionRnToR(traction[i][0]);
			FunctionRnToR ldy = new ConstantFunctionRnToR(traction[i][1]);
			Node nt1 = mesh.getNode(edge[0] - 1);
			Node nt2 = mesh.getNode(edge[1] - 1);
			Node nt3 = mesh.getNode(edge[2] - 1);
			Edge e1 = mesh.addEdge(new Edge(nt1, nt2));
			Load el1 = new DistributedLoad(e1, ldx, ldy);
			m.addLoad(el1);
			Edge e2 = mesh.addEdge(new Edge(nt2, nt3));
			Load el2 = new DistributedLoad(e2, ldx, ldy);
			m.addLoad(el2);
		}

		return m;
	}

	// @Test
	// public void testKt() {
	// Model m = new volumetric_locking_demo_quadratic().createModel(0);
	// m.getMesh().setNumberOfFields(2);
	// m.preAnalysis();
	// int n = m.getMesh().countDofs();
	// IMatrix kt = new Array2DMatrix(n, n);
	// m.assembleKt(kt);
	// System.out.println(MatrixFormat.format(kt));
	// }

	@Test
	public void testSolve1() throws SolveFailedException {
		Model M = createModel(0);

		for (int i = 0; i < M.countElements(); i++) {
			((StructuralElement) M.getElement(i))
					.setIntegrationType(IntegrationType.REDUCED_INTEGRATION);
		}

		solveAndCheckWithTolerance(M, 1e-1);
	}

	@Test
	public void testSolve2() throws SolveFailedException {
		Model M = createModel(0);

		for (int i = 0; i < M.countElements(); i++) {
			((StructuralElement) M.getElement(i))
					.setIntegrationType(IntegrationType.SELECTIVE_REDUCED_INTEGRATION);
		}

		solveAndCheckWithTolerance(M, 1e-1);
	}

	@Test
	public void testSolve3() throws SolveFailedException {
		Model M = createModel(1);

		solveAndCheckWithTolerance(M, 1e-1);
	}

	@Test
	public void testSolve4() throws SolveFailedException {
		Model M = createModel(2);

		solveAndCheckWithTolerance(M, 1);
	}

	@Test
	public void testSolve5() throws SolveFailedException {
		Model M = createModel(3);

		solveAndCheckWithTolerance(M, 1);
	}

	private double[][] u_ana = new double[][] { { 0, 0.5334 }, { 0.1066373504, 0.5225230169 },
			{ 0.2045738463, 0.495284049 }, { 0.2979591953, 0.4416180931 },
			{ 0.3756341878, 0.3756341878 }, { 0.4416180931, 0.2979591953 },
			{ 0.495284049, 0.2045738463 }, { 0.5225230169, 0.1066373504 }, { 0.5334, 0 },
			{ 0, 0.3048785714 }, { 0.1163149082, 0.2812390318 }, { 0.2151364301, 0.2151364301 },
			{ 0.2812390318, 0.1163149082 }, { 0.3048785714, 0 }, { 0, 0.2135 },
			{ 0.0418955657, 0.2094778284 }, { 0.0818831987, 0.197031447 },
			{ 0.118544905, 0.1773909369 }, { 0.1507771337, 0.1507771337 },
			{ 0.1773909369, 0.118544905 }, { 0.197031447, 0.0818831987 },
			{ 0.2094778284, 0.0418955657 }, { 0.2135, 0 }, { 0, 0.1643192308 },
			{ 0.0628420553, 0.1520372305 }, { 0.1160953623, 0.1160953623 },
			{ 0.1520372305, 0.0628420553 }, { 0.1643192308, 0 }, { 0, 0.1336 },
			{ 0.0260930146, 0.1311341246 }, { 0.0510036961, 0.1233422716 },
			{ 0.0740680739, 0.1111021109 }, { 0.0944171708, 0.0944171708 },
			{ 0.1111021109, 0.0740680739 }, { 0.1233422716, 0.0510036961 },
			{ 0.1311341246, 0.0260930146 }, { 0.1336, 0 }, };

	private void solveAndCheckWithTolerance(Model M, double tol) throws SolveFailedException {
		Mesh mesh = M.getMesh();
		/*
		 * analysis
		 */
		Solver solver = new BaitschSolver(SolverType.BAITSCH_SOLVER_SPOOLES);
		mesh.setNumberOfFields(2);
		M.preAnalysis();
		M.Initialize();
		int n = mesh.countDofs();
		solver.setSize(n);
		M.assembleKt(solver.getA());
		double[] r = new double[n];
		r = M.assembleR(r, 0.0);
		double[] u = r.clone();
		solver.solve(u);
		M.setX(1, u);
		M.FinalizeNonlinearIteration();
		// StructuralElement se = (StructuralElement) M.getElement(0);
		// System.out.println(MatrixFormat.format(se.getSigma(0)));

		double[] ri = new double[n];
		M.assembleRi(ri);
		System.out.println("norm r = " + LinAlg.norm2(n, r));
		System.out.println("norm ri = " + LinAlg.norm2(n, ri));
		System.out.println("residual = " + LinAlg.diffNorm2(n, r, ri));

		/*
		 * check with correct result
		 */
		for (int i = 0; i < mesh.countNodes(); i++) {
			double[] uHatLocal = mesh.getNode(i).getUHatLocal();
			Assert.assertEquals(u_ana[i][0], uHatLocal[0], tol);
			Assert.assertEquals(u_ana[i][1], uHatLocal[1], tol);
		}
	}

	@SuppressWarnings("unused")
	private void seeGaussPoints() {
		Model M = createModel(0);
		/*
		 * check Gauss point
		 */
		Element e = M.getElement(0);
		int ng = e.getNumIntegrationPoints();
		for (int i = 0; i < ng; i++) {
			double[] xi = e.getIntegrationPoint(i);
			System.out.println("Gauss point " + i + ": xi[0] = " + xi[0] + ", xi[1] = " + xi[1]);
		}
	}

	public static void main(String[] args) throws SolveFailedException, IOException {
		volumetric_locking_demo_quadratic problem = new volumetric_locking_demo_quadratic();
		Model M = problem.createModel(0);

		for (int i = 0; i < M.countElements(); i++) {
			// ((StructuralElement) M.getElement(i))
			// .setIntegrationType(IntegrationType.SELECTIVE_REDUCED_INTEGRATION);
			// //worse than reduced integration

			((StructuralElement) M.getElement(i))
					.setIntegrationType(IntegrationType.REDUCED_INTEGRATION); // best

			// ((StructuralElement) M.getElement(i)).enableAndInitializeBBar();
			// // best
		}

		/*
		 * analysis
		 */
		M.getMesh().setNumberOfFields(2);
		Analysis an = new LinearStaticAnalysis(M);
		an.getStrategy().setEchoLevel(EchoLevelType.OUTPUT);
		an.getStrategy().addObserver(new StressObserver(M));
		an.run();

		/*
		 * post process
		 */
		double[][] error = new double[M.getMesh().countNodes()][2];
		for (int i = 0; i < M.getMesh().countNodes(); i++) {
			double[] uHatLocal = M.getMesh().getNode(i).getUHatLocal();
			System.out.println("node " + (i + 1) + ":" + ArrayFormat.format(uHatLocal));
			error[i][0] = Math.abs(uHatLocal[0] - problem.u_ana[i][0]);
			error[i][1] = Math.abs(uHatLocal[1] - problem.u_ana[i][1]);
		}

		String projDir = "/home/hbui/workspace2/src-matlab/locking";
		String projName = "volumetric_locking_demo_quadratic.gid";
		GidPostStructural post = new GidPostStructural(projDir, projName);
		post.clearAllPostFiles();
		post.writeMesh(an.getModel(), 0);
		post.writeNodalDisplacementResults(M, 0);
		post.writeGaussPointStressResults(M, 0, StressType.CAUCHY_STRESS);

		post.writeNodalVectorResult("EXACT_DISPLACEMENT", problem.u_ana, 0);
		post.writeNodalVectorResult("DISPLACEMENT_ERROR", error, 0);
	}
}
